import os
import numpy as np
from torch.utils.data import Dataset
import SimpleITK as sitk
import torch.nn.functional as F
import random
import matplotlib.pyplot as plt
from scipy.ndimage.interpolation import zoom

def crop_from_centre(centre_x,centre_y,crop_size):

    crop_x = crop_size[0] / 2
    crop_y = crop_size[1] / 2

    left = down = right = up = front = 0
    back = 24
    if crop_x <= centre_x <= 256 - crop_x and crop_y <= centre_y <= 256 - crop_y:
        left = centre_x - crop_x
        right = centre_x + crop_x
        down = centre_y - crop_y
        up = centre_y + crop_y

    elif centre_x <= crop_x and crop_y <= centre_y <= 256 - crop_y:  #
        left = 0
        right = crop_size[0]
        down = centre_y - crop_y
        up = centre_y + crop_y
    elif centre_x >= 256 - crop_x and crop_y <= centre_y <= 256 - crop_y:  #
        left = 255-crop_size[0]
        right = 255
        down = centre_y - crop_y
        up = centre_y + crop_y

    elif crop_x <= centre_x <= 256 - crop_x and centre_y <= crop_y:  #
        left = centre_x - crop_x
        right = centre_x + crop_x
        down = 0
        up = crop_size[1]
    elif crop_x <= centre_x <= 256 - crop_x and centre_y >= 256 - crop_y:  #
        left = centre_x - crop_x
        right = centre_x + crop_x
        down = 255-crop_size[1]
        up = 255

    coordinates = (left, down, front, right, up, back)
    print('    coordinates ', coordinates)
    return coordinates

def patch_generator(tumor_label):
    assert tumor_label.ndim == 3
    # tumor_label = (tumor_label > 0).astype(np.int)
    label = tumor_label
    annotated_digits = np.nonzero(label)  # get annotated voxel indices
    # tumor_digits = np.nonzero(tumor_label)
    assert np.sum(label) > 1, annotated_digits

    min_vertex = np.array([np.min(annotated_digits[0]), np.min(annotated_digits[1]), np.min(annotated_digits[2])])
    max_vertex = np.array([np.max(annotated_digits[0]), np.max(annotated_digits[1]), np.max(annotated_digits[2])])
    print('annotated_digits', min_vertex, max_vertex)

    centre_x = (np.min(annotated_digits[0])+np.max(annotated_digits[0])) //2
    centre_y = (np.min(annotated_digits[1])+np.max(annotated_digits[1])) //2
    print('             centre :',centre_x,centre_y)


    x_len = np.max(annotated_digits[0])-np.min(annotated_digits[0])
    y_len = np.max(annotated_digits[1]) - np.min(annotated_digits[1])

    print(f"x_len, y_len is {x_len, y_len}")
    print("*" * 100)
    if x_len <=60 and y_len<=60:
        coordinates = crop_from_centre(centre_x,centre_y,crop_size=(75,75))
    else:
        coordinates = crop_from_centre(centre_x,centre_y,crop_size=(100,100))
    return coordinates

def crop_data_by_coordinates(three_modalities_arr, coordinates):
    bboxes = []

    print(coordinates)
    for arr in three_modalities_arr:
        bbox = arr[int(coordinates[0]): int(coordinates[3]), int(coordinates[1]):int(coordinates[4]),
             int(coordinates[2]): int(coordinates[5])]
        bboxes.append(bbox)

    return bboxes

def load_data(t1_path,t2_path,mask_path): # 导入path
    # get data and label
    t1_img = sitk.ReadImage(t1_path)
    t1_img = sitk.GetArrayFromImage(t1_img)
    t1_img = t1_img.transpose((1, 2, 0)).astype(np.int)

    t2_img = sitk.ReadImage(t2_path)
    t2_img = sitk.GetArrayFromImage(t2_img)
    t2_img = t2_img.transpose((1, 2, 0)).astype(np.int)

    mask = sitk.ReadImage(mask_path)
    mask = sitk.GetArrayFromImage(mask)
    tumor_label = mask.transpose((1, 2, 0)).astype(np.int)

    coordinates = patch_generator(tumor_label)
    assert len(coordinates) == 6, "unacceptable coordinates for cropping"
    # print('size big',coordinates[3]-coordinates[0] , coordinates[4]-coordinates[1] , coordinates[5]-coordinates[2] )
    t1_bbox = t1_img[int(coordinates[0]): int(coordinates[3]), int(coordinates[1]):int(coordinates[4]),
             int(coordinates[2]): int(coordinates[5])]
    t2_bbox = t2_img[int(coordinates[0]): int(coordinates[3]), int(coordinates[1]):int(coordinates[4]),
            int(coordinates[2]): int(coordinates[5])]

    # 线性插值
    if t1_bbox.shape[0] < 100:
        new_shape = (100, 100, 24)
        resize_factor = np.array(new_shape) / np.array(t1_bbox.shape)
        t1_bbox = zoom(t1_bbox, resize_factor, mode='nearest')
        t2_bbox = zoom(t2_bbox, resize_factor, mode='nearest')

    return t1_bbox, t2_bbox


class raw_Dataset(Dataset):

    def __init__(self,  subset, inference=False):
        super(raw_Dataset, self).__init__()
        t1_image = []
        t2_image = []
        mask = []
        patients = []
        # 我自己整合划分后的数据集
        source_path = "/mnt/hard_disk/liuwennan/dataset/meningioma_data/Data_processed/bbox/tumor/label/"
        print('source path:',source_path)
        for subtype in os.listdir(source_path):
            subtype_path = os.path.join(source_path, str(subtype))
            for patient in os.listdir(subtype_path):
                patient_path = os.path.join(subtype_path, patient)
                patients.append(patient_path)
                for series_IDs in os.listdir(patient_path):
                    if series_IDs =='t1.nii.gz':
                        t1_path = os.path.join(patient_path,series_IDs)
                        t1_image.append(t1_path)
                    elif series_IDs =='t2.nii.gz':
                        t2_path = os.path.join(patient_path,series_IDs)
                        t2_image.append(t2_path)
                    elif series_IDs =='t1_mask.nii.gz':  # label data:t1_mask   unlabel data:t2_mask
                        mask_path = os.path.join(patient_path,series_IDs)
                        mask.append(mask_path)

        self.name = patients
        self.t1_img = t1_image
        self.t2_img = t2_image
        self.mask = mask


    def __getitem__(self, index):

        # get image and label
        print('tobe deal with :',self.name[index])
        t1_bbox, t2_bbox= load_data(self.t1_img[index],self.t2_img[index],self.mask[index])
        name_path = self.name[index]
        t1_bbox = t1_bbox.transpose((2,0,1)).astype(np.int)
        t2_bbox = t2_bbox.transpose((2,0,1)).astype(np.int)
        print('---- bbox_shape',t1_bbox.shape)
        return t1_bbox, t2_bbox, name_path

    def __len__(self):
        return len(self.t1_img)


if __name__ == '__main__':


    origin = (0.0, 0.0, 0.0)
    spacing = (1.0, 1.0, 1.0)
    direction = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0)
    # train_set = Label_Dataset('label', inference=False)
    train_set = raw_Dataset('label', inference=False)
    for i in range(len(train_set)):
        t1_bbox, t2_bbox,name_path = train_set[i]
        t1_bbox = t1_bbox.astype(np.float32)
        t2_bbox = t2_bbox.astype(np.float32)

        t1_bbox_nii = sitk.GetImageFromArray(t1_bbox)
        t1_bbox_nii.SetOrigin(origin)
        t1_bbox_nii.SetSpacing(spacing)
        t1_bbox_nii.SetDirection(direction)
        sitk.WriteImage(t1_bbox_nii, name_path + '/' + 't1_bbox.nii.gz')
        print('save:', name_path + '/' + 't1_bbox.nii.gz')

        t2_bbox_nii = sitk.GetImageFromArray(t2_bbox)
        t2_bbox_nii.SetOrigin(origin)
        t2_bbox_nii.SetSpacing(spacing)
        t2_bbox_nii.SetDirection(direction)
        sitk.WriteImage(t2_bbox_nii, name_path + '/' + 't2_bbox.nii.gz')
        print('save:', name_path + '/' + 't2_bbox.nii.gz')



